package com.gzx.asr.client;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.gzx.asr.config.FunAsrConfig;
import com.gzx.asr.constant.ModeConstant;
import com.gzx.asr.modular.text.vo.GzxAsrVoiceTextVo;
import lombok.extern.slf4j.Slf4j;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;

import java.io.File;
import java.io.FileInputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

@Slf4j
public class FunAsrWsClient extends WebSocketClient {

    //临时存储ws的结果
    private final List<GzxAsrVoiceTextVo> results = new CopyOnWriteArrayList<>();
    private final CountDownLatch latch = new CountDownLatch(1);

    private boolean iseof = false;
    String strChunkSize = "5,10,5";
    int chunkInterval = 10;
    int sendChunkSize = 3200;

    String wavPath;
    String mode;
    String wavName;
    String hotWords;

    public FunAsrWsClient(URI serverURI, String mode, String wavPath, String wavName, String hotWords) {
        super(serverURI);
        this.mode = mode;
        this.wavPath = wavPath;
        this.wavName = wavName;
        this.hotWords = hotWords;
    }

    @Override
    public void onOpen(ServerHandshake handshakedata) {
        recWav();//连接创建后的回调函数
    }

    @Override
    public void onMessage(String message) {
        JSONObject resp = new JSONObject();
        try {
            resp = JSON.parseObject(message);
            log.info("resp: {}", resp);
            GzxAsrVoiceTextVo vo = new GzxAsrVoiceTextVo();
            // 解析带说话人标签的文本
            if (resp.containsKey("text")) {
                String text = resp.getString("text");
                if (text.contains("<spk")) {
                    //vo.setSpeakerRole(0);
                    text = text.replace("<spk0>", "[客服]")
                            .replace("<spk1>", "[客户]");
                }
                //保存每一条消息
                vo.setContent(text);
            }

            if (resp.containsKey("timestamp")) {
                vo.setChatTime(resp.getString("timestamp"));
            }
            results.add(vo);
        } catch (Exception e) {
            log.error("接收服务器消息异常", e);
        }
        if (iseof && ModeConstant.OFF_LINE_MODE.equals(mode) && resp.containsKey(ModeConstant.FINAL_JSONKEY) && resp.getBoolean(ModeConstant.FINAL_JSONKEY)) {
            log.info("识别结束");
            close();
        }
    }

    @Override
    public void onClose(int code, String reason, boolean remote) {
        latch.countDown();
        log.info("Connection closed by {} Code: {} Reason: {}", remote ? "remote peer" : "us", code, reason);
    }

    @Override
    public void onError(Exception ex) {
        log.error("ex:", ex);
        latch.countDown();
        // if the error is fatal then onClose will be called additionally
    }

    // 提供同步获取结果的方法
    public List<GzxAsrVoiceTextVo> getResults() throws InterruptedException {
        latch.await(); // 等待连接关闭
        return new ArrayList<>(results); // 返回不可变副本
    }

    // 提供带超时的获取方式（可选）
    public List<GzxAsrVoiceTextVo> getResults(long timeout, TimeUnit unit) throws InterruptedException, TimeoutException {
        boolean completed = latch.await(timeout, unit);
        if (!completed) {
            throw new TimeoutException("WebSocket 连接超时");
        }
        return new ArrayList<>(results);
    }

    // 发送asr配置
    public void sendConfig(String mode, String strChunkSize, int chunkInterval, String wavName, boolean isSpeaking, String suffix, String hotwords) {
        JSONObject initConfig = FunAsrConfig.initConfig(mode, strChunkSize, chunkInterval, wavName, isSpeaking, suffix, hotwords);
        if (initConfig != null) {
            send(initConfig.toString());
        }
    }

    // 发送结束指令
    public void sendEof() {
        try {
            JSONObject obj = new JSONObject();
            obj.put("is_speaking", false);
            log.info("发送sendEof指令:{}", obj);
            send(obj.toString());
            iseof = true;
        } catch (Exception e) {
            log.error("sendEof指令异常", e);
        }
    }

    // 读取录音文件，向服务器发送录音内容
    public void recWav() {
        String fileName = wavPath;
        String suffix = fileName.split("\\.")[fileName.split("\\.").length - 1];
        sendConfig(mode, strChunkSize, chunkInterval, wavName, true, suffix, hotWords);
        File file = new File(wavPath);
        //因为音频文件导致没有生成左右声道文件
        if (!file.exists()) {
            close();
            return;
        }

        int chunkSize = sendChunkSize;
        byte[] bytes = new byte[chunkSize];

        int readSize;
        try (FileInputStream fis = new FileInputStream(file)) {
            if (wavPath.endsWith(".wav")) {
                fis.read(bytes, 0, 44); //skip first 44 wav header
            }
            readSize = fis.read(bytes, 0, chunkSize);
            while (readSize > 0) {
                byte[] audioChunk = Arrays.copyOf(bytes, readSize);
                // send when it is chunk size
                if (readSize == chunkSize) {
                    send(audioChunk); // send buf to server

                } else {
                    send(Arrays.copyOf(audioChunk, readSize));
                }
                // if not in offline mode, we simulate online stream by sleep
                if (!ModeConstant.OFF_LINE_MODE.equals(mode)) {
                    Thread.sleep(chunkSize / 32);
                }
                readSize = fis.read(bytes, 0, chunkSize);
            }

            if (!ModeConstant.OFF_LINE_MODE.equals(mode)) {
                // if not offline, we send eof and wait for 3 seconds to close
                Thread.sleep(2000);
                sendEof();
                Thread.sleep(3000);
                close();
            } else {
                // if offline, just send eof
                sendEof();
            }
        } catch (InterruptedException ex) {
            Thread.currentThread().interrupt();
        } catch (Exception e) {
            log.error("向服务器发送消息异常", e);
        }
    }

    public static void main(String[] args) throws URISyntaxException, InterruptedException {
        //单线程方式
        String wavPath = "D:\\asr\\download\\202505021014071406085700501746152047.wav";
        String mode = "offline";
        String wavName = "202505021014071406085700501746152047.wav";
        String hotWords = "";
        String srvIp = "117.72.49.18";
        String srvPort = "10096";
        String wsAddress = "ws://" + srvIp + ":" + srvPort;
        FunAsrWsClient c = new FunAsrWsClient(new URI(wsAddress), mode, wavPath, wavName, hotWords);
        c.connect();
    }

}
